from __future__ import division

import collections
import math
import operator
from pyre.core import types

_inf = float('inf')

def expand_adjacent(node, world):
    """
    expand_adjacent(Coordinate, World): generator of Coordinate
    Yields all the passable positions adjacent to the given position
    """

    for dx in [-1, 0, 1]:
        for dy in [-1, 0, 1]:
            pos = node + (dx, dy)
            if world.passable_at(pos):
                yield pos


class DijkstraMap(object):

    def __init__(self, world, grid):
        self.world = world
        self.grid = grid

    @classmethod
    def compute(cls, world, goals):
        dmap = [[_inf for _ in range(world.rows)] for _ in range(world.cols)]
        frontier = collections.deque()
        explored = set()

        for goal in goals:
            dmap[goal[0]][goal[1]] = 0
            frontier.append(types.Coordinate.copy(goal))

        while frontier:
            curr = frontier.popleft()
            if curr not in explored:
                explored.add(curr)

                proposed_cost = dmap[curr[0]][curr[1]] + 1
                for node in expand_adjacent(curr, world):
                    old_cost = dmap[node[0]][node[1]]
                    if proposed_cost < old_cost:
                        dmap[node[0]][node[1]] = proposed_cost
                        frontier.append(node)

        return cls(world, dmap)

    @classmethod
    def compute_safety(cls, world, hazards):
        dmap = cls.compute(world, hazards)
        dmap *= -1.2
        dmap._iterative_update()
        return dmap

    def _iterative_update(self):
        changed = True
        while changed:
            changed = False
            for x in range(1, self.world.cols - 1):
                for y in range(1, self.world.rows - 1):
                    if world.passable_at((x, y)):
                        for dx in [-1, 0, 1]:
                            for dy in [-1, 0, 1]:
                                if self[x][y] > self[x + dx][y + dy] + 1:
                                    self[x][y] = self[x + dx][y + dy] + 1
                                    changed = True

    def __getitem__(self, index):
        return self.grid[index]

    def _copy(self):
        dmap = [[self[x][y] for y in range(self.world.rows)] 
                for x in range(self.world.cols)]
        return DijkstraMap(self.world, dmap)

    def _scale(self, scalar):
        for x in range(self.world.cols):
            for y in range(self.world.rows):
                if self.world.passable_at((x, y)):
                    self[x][y] *= scalar

    def _blend(self, other, func=operator.add):
        assert self.world is other.world

        for x in range(self.world.cols):
            for y in range(self.world.rows):
                if self.world.passable_at((x, y)):
                    self[x][y] = func(self[x][y], other[x][y])

    def __add__(self, other):
        result = self._copy()
        result._blend(other)
        return result

    def __iadd__(self, other):
        self._blend(other)
        return self

    def __sub__(self, other):
        result = self._copy()
        result._blend(other, operator.sub)
        return result

    def __iadd__(self, other):
        self._blend(other, operator.sub)
        return self

    def __mul__(self, scalar):
        result = self._copy()
        result._scale(scalar)
        return result

    def __imul__(self, scalar):
        self._scale(scalar)
        return self

    def __div__(self, inv_scalar):
        return self.__mul__(1 / inv_scalar)

    def __idiv__(self, inv_scalar):
        return self.__imul__(1 / inv_scalar)
